import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.metrics import r2_score, mean_squared_error
import matplotlib.pyplot as plt
import seaborn as sns


# 1. 数据准备
def prepare_data(file_path):
    # 读取xlsx文件
    df = pd.read_excel(file_path)

    # 设置中文字体
    plt.rcParams['font.sans-serif'] = ['SimHei']
    plt.rcParams['axes.unicode_minus'] = False

    # 定义自变量(X)和因变量(y)
    X = df[['GDP（万亿）', 'GDP增长率（%）', '第一产业GDP（万亿元）',
            '第二产业GDP（万亿元）', '第三产业GDP（万亿元）',
            '中国区域创新能力指数（基准年=2004）', '恩格尔系数（%）']]
    y = df['就业率（估算，%）']

    # 创建简化的列名便于后续处理
    column_mapping = {
        'GDP（万亿）': 'GDP',
        'GDP增长率（%）': 'GDP增长率',
        '第一产业GDP（万亿元）': '第一产业GDP',
        '第二产业GDP（万亿元）': '第二产业GDP',
        '第三产业GDP（万亿元）': '第三产业GDP',
        '中国区域创新能力指数（基准年=2004）': '创新能力指数',
        '恩格尔系数（%）': '恩格尔系数'
    }

    X.columns = [column_mapping[col] for col in X.columns]

    return X, y, df['年份']


# 2. 构建模型
def build_model(X, y):
    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    # 创建并训练模型
    model = LinearRegression()
    model.fit(X_train, y_train)

    # 模型评估
    y_pred = model.predict(X_test)
    r2 = r2_score(y_test, y_pred)
    mse = mean_squared_error(y_test, y_pred)
    rmse = np.sqrt(mse)

    # 获取系数和截距
    coefficients = pd.DataFrame({
        '变量': X.columns,
        '系数': model.coef_
    })
    intercept = model.intercept_

    return model, coefficients, intercept, r2, rmse, X_test, y_test, y_pred


# 3. 模型分析和可视化
def analyze_model(X, y, years, coefficients, r2, rmse):
    # 创建图表目录
    import os
    if not os.path.exists('model_analysis'):
        os.makedirs('model_analysis')

    # 1. 相关性分析
    correlation_matrix = pd.concat([X, y], axis=1).corr()
    plt.figure(figsize=(12, 10))
    sns.heatmap(correlation_matrix, annot=True, cmap='coolwarm', center=0, fmt='.2f')
    plt.title('变量相关性热力图')
    plt.tight_layout()
    plt.savefig('model_analysis/correlation_heatmap.png')
    plt.close()

    # 2. 各变量对就业率的影响系数可视化
    plt.figure(figsize=(12, 6))
    bars = plt.bar(coefficients['变量'], coefficients['系数'])
    plt.xticks(rotation=45, ha='right')
    plt.title('各变量对就业率的影响系数')
    plt.ylabel('系数值')

    # 为每个柱状图添加数值标签
    for bar in bars:
        height = bar.get_height()
        plt.text(bar.get_x() + bar.get_width() / 2., height,
                 f'{height:.4f}',
                 ha='center', va='bottom')

    plt.tight_layout()
    plt.savefig('model_analysis/coefficients_impact.png')
    plt.close()

    # 3. 就业率随时间变化趋势
    plt.figure(figsize=(12, 6))
    plt.plot(years, y, 'b-', label='实际就业率')
    plt.xlabel('年份')
    plt.ylabel('就业率(%)')
    plt.title('就业率历年变化趋势(1984-2023)')
    plt.xticks(rotation=45)
    plt.grid(True)
    plt.legend()
    plt.tight_layout()
    plt.savefig('model_analysis/employment_trend.png')
    plt.close()

    # 4. 生成分析报告
    with open('model_analysis/analysis_report.txt', 'w', encoding='utf-8') as f:
        f.write('就业率影响因素分析报告\n')
        f.write('=' * 50 + '\n\n')

        f.write('1. 模型性能指标：\n')
        f.write(f'   R² 得分: {r2:.4f}\n')
        f.write(f'   均方根误差(RMSE): {rmse:.4f}\n\n')

        f.write('2. 各因素影响程度（系数）：\n')
        for _, row in coefficients.iterrows():
            f.write(f'   {row["变量"]}: {row["系数"]:.4f}\n')

        f.write('\n3. 变量相关性分析：\n')
        f.write(correlation_matrix.to_string())

    return correlation_matrix


# 主函数
def main(file_path):
    # 准备数据
    X, y, years = prepare_data(file_path)

    # 构建模型
    model, coefficients, intercept, r2, rmse, X_test, y_test, y_pred = build_model(X, y)

    # 分析模型并生成可视化
    correlation_matrix = analyze_model(X, y, years, coefficients, r2, rmse)

    # 打印模型方程
    print("\n就业率预测模型方程:")
    print(f"就业率 = {intercept:.4f}", end="")
    for i, row in coefficients.iterrows():
        print(f" + ({row['系数']:.4f} × {row['变量']})", end="")
    print("\n")

    print("模型评估:")
    print(f"R² 得分: {r2:.4f}")
    print(f"均方根误差(RMSE): {rmse:.4f}")
    print("\n各变量影响系数:")
    print(coefficients)

    print("\n分析报告和可视化结果已保存到 'model_analysis' 文件夹")

    return model, coefficients, correlation_matrix


# 运行模型
file_path = './data/GDP.xlsx'  # 替换为实际的文件路径
model, coefficients, correlation_matrix = main(file_path)